import javax.swing.*;
import org.math.plot.*;
import static org.math.array.DoubleArray.*;
import static org.math.array.LinearAlgebra.*;
import org.math.array.util.*;
import java.io.*;
import static java.lang.Math.*;

/**
 * Practical1
 * Use the idea of least-squares fitting to fit functions to data
 * 
 * @author Richard Wilson
 * @version 1
 */

public class Practical1
{
    // Preamble
    private BufferedReader stdin;

    
    public Practical1()
    {
        stdin = new BufferedReader(new InputStreamReader(System.in));
    }
    
    
    /***************************** PART A ******************************/
    /**
     * Implements part A of the practical
     * Complete the function simpleLeastSquares and then run partA() to see the results
     */
    public void partA()
    {
        double [][] Dataset= { {1,1.1}, {2,2.05}, {3,2.9}, {4,4.1}, {5,5.02} };
        double [][] p;
             
        // Complete this method below
        p=simpleLeastSquares(Dataset);
        
        // Generate data for our fitted line
        double[] FittedX={ 0,3,6 };
        int n=FittedX.length;
        double[] FittedY=new double[n];
        for( int i=0;i<n;i++ ) {
            FittedY[i]=simpleLinearFunction(p,FittedX[i]);
        }
        
        // Plot the fitting result
        plot2D(Dataset,FittedX,FittedY);
        
        
    }
    
     /**
     * Least-squares fit of y=ax+b
     * 
     * @param  X   The dataset
     * @return x   =[a; b] the parameters 
     */
    private double[][] simpleLeastSquares(double[][] X)
    {
        int n=X.length;
        
      
        //-------------------------------------------------//
        // Calculate the values of matrices A and c
        
        double sqx=0;
        double x=0;
        double y=0;
        double xy=0;
        
        for(int i=0; i<n; i++) {
            sqx+=X[i][0]*X[i][0];
            x+=X[i][0];
            y+=X[i][1];
            xy+=X[i][0]*X[i][1];
        }
        
        double[][] A={{sqx,x} , {x,n}};
        double[][] c={ {xy},{y}};
        double[][] invA=inverse(A);
        
   
        double[][] p=times(invA,c);
        
                for(int i=0; i<p.length; i++) {
            for(int j=0; j<p[0].length; j++) {
                System.out.println(p[i][j] + " ");
            }
        }
        
        //-------------------------------------------------//

        return p;
    }

    
    /**
     * Return the value of the linear function y=ax+b
     * 
     * @param  p   =[a; b] the parameters 
     * @param  x   The value of x
     * @return y   The value of the function
     */
    private double simpleLinearFunction(double[][] p,double x)
    {
        return p[0][0]*x+p[1][0];
    }
        


    /***************************** PART B ******************************/
    /**
     * Implements part B of the practical
     * Complete the function threeDLeastSquares and then run partB() to see the results
     */
        
    public void partB()
    {
        double [][] Dataset= { {1,1,0.5}, {2,2,3}, {3,1,4}, {1,3,4.5}, {3,3,5.5}, {4,2,6}, {2,4,7} };
        double [][] p;
        
        // Complete this method below
        p=threeDLeastSquares(Dataset);
        
        // Generate data for our fitted line
        // This must be on a grid pattern for the plotter
        double[] FittedX={ 0,2.5,5 };
        double[] FittedY={ 0,2.5,5 };
        double[][] FittedZ=new double[3][3];
        int i,j;
        for( i=0;i<3;i++ ) {
            for( j=0;j<3;j++ ) {
                FittedZ[i][j]=threeDLinearFunction(p,FittedX[i],FittedY[j]);
            }   
        }
        
        // Plot the fitting result
        plot3D(Dataset,FittedX,FittedY,FittedZ);        
        
    }
    
     /**
     * Least-squares fit of z=ax+by+c
     * 
     * @param  X   The dataset
     * @return x   =[a; b; c] the parameters 
     */
    private double[][] threeDLeastSquares(double[][] X)
    {
        int n=X.length;
        
        // Calculate the matrix A and vector c
        // You will need to work out the correct formulas
        // We store everything as a matrix for consistency
        //-------------------------------------------------//
        // Calculate the values of matrices A and c
        
        double sqx=0;
        double x=0;
        double y=0;
        double xy=0;
        double sqy=0;
        double xz=0;
        double yz=0;
        double z=0;
        
        for(int i=0; i<n; i++) {
            sqx+=X[i][0]*X[i][0];
            x+=X[i][0];
            y+=X[i][1];
            xy+=X[i][0]*X[i][1];
            sqy+=X[i][1]*X[i][1];
            xz+=X[i][0]*X[i][2];
            yz+=X[i][1]*X[i][2];
            z+=X[i][2];
        }
        
        double[][] A={{sqx,sqy,x} , {xy,sqy,y} ,{x,y,n}};
        double[][] c={ {xz},{yz},{z}};
        double[][] invA=inverse(A);
        
   
        double[][] p=times(invA,c);
        
        //                 for(int i=0; i<p.length; i++) {
        //             for(int j=0; j<p[0].length; j++) {
        //                 System.out.println(p[i][j] + " ");
        //             }
        //         }
                

        
        //-------------------------------------------------//

        return p;
    }
    /**
     * Return the value of the linear function y=ax+b
     * 
     * @param  p   =[a; b] the parameters 
     * @param  x   The value of x
     * @return y   The value of the function
     */
    private double threeDLinearFunction(double[][] p,double x,double y)
    {
        return p[0][0]*x+p[1][0]*y+p[2][0];
    }

    
    
    /***************************** PART C ******************************/
    /**
     * Implements part C of the practical
     * Complete the function complexLeastSquares and then run partC() to see the results
     */
    public void partC()
    {
        double [][] Dataset= { {0,3.1}, {1,5.8}, {2,7.1}, {3, 6.1}, {4,2.7}, {5,-2.0} };
        double [][] p;
        
        // Complete this method below
        p=complexLeastSquares(Dataset);
        
        // Generate data for our fitted line
        int n=20;
        double[] FittedX=new double[n];
        double[] FittedY=new double[n];
        for( int i=0;i<n;i++ ) {
            FittedX[i]=i*6.0/(n-1);
            FittedY[i]=nonLinearFunction(p,FittedX[i]);
        }
        
        // Plot the fitting result
        plot2D(Dataset,FittedX,FittedY);
        
        
    }
    
    /**
     * Least-squares fit of z=ax^2+bx+c
     * 
     * @param  X   The dataset
     * @return x   =[a; b; c] the parameters 
     */
    private double[][] complexLeastSquares(double[][] X)
    {
        int n=X.length;
        
        // Calculate the matrix A and vector c
        // You will need to work out the correct formulas
        // We store everything as a matrix for consistency
       //-------------------------------------------------//
        // Calculate the values of matrices A and c

        double sqx=0;
        double x=0;
        double y=0;
        double xy=0;
        double quadx=0;
        double cubex=0;
        double xxy=0;
        
        for(int i=0; i<n; i++) {
            sqx+=X[i][0]*X[i][0];
            x+=X[i][0];
            y+=X[i][1];
            xy+=X[i][0]*X[i][1];
            cubex+=(X[i][0])*(X[i][0])*(X[i][0]);
            quadx+=(X[i][0])*(X[i][0])*(X[i][0])*(X[i][0]);
            xxy+=(X[i][0])*(X[i][0])*(X[i][1]);
        }
        
        double[][] A={{quadx,cubex,sqx} , {cubex,sqx,x}, {sqx, x, n}};
        double[][] c={ {xxy},{xy},{y}};
        double[][] invA=inverse(A);
        
   
        double[][] p=times(invA,c);
        
                for(int i=0; i<p.length; i++) {
            for(int j=0; j<p[0].length; j++) {
                System.out.println(p[i][j] + " ");
            }
        }

        //-------------------------------------------------//

        return p;
    }

    /**
     * Return the value of the nonlinear function y=ax^2+b*x+c
     * 
     * @param  p   =[a; b] the parameters 
     * @param  x   The value of x
     * @return y   The value of the function
     */
    private double nonLinearFunction(double[][] p,double x)
    {
        return p[0][0]*x*x+p[1][0]*x+p[2][0];
    }
    
    
    
    

    
    
    
    /**    Utility Routines       */
    
    /**
     * Display the contents of an array and wait for Enter key
     * 
     * @param  info initial message
     * @param  M    Array
     */
    private void showArrayAndWait(String info,double[][] M)
    {
        System.out.println(info+"\n"+org.math.array.DoubleArray.toString(M));
        System.out.println("Press Enter key...");
        try {
            stdin.readLine();
        } catch (IOException ex) {
        }
    }
    
    
    /**
     * Plot points and fit in 2D
     * 
     * @param  D   Datapoints which were fitted
     * @param  fx  fitted line x values
     * @param  fy  fitted line y values
     */
    private void plot2D(double[][] D,double[] fx,double[] fy)
    {
        Plot2DPanel plot = new Plot2DPanel();
        plot.addScatterPlot("Datapoints", getColumnCopy(D,0), getColumnCopy(D,1));
        plot.addLinePlot("Fitted line", fx, fy);
        JFrame frame = new JFrame("2D Fit Plot");
        frame.setSize(600, 600);
        frame.setContentPane(plot);
        frame.setVisible(true);
    }
    
    /**
     * Plot points and fit in 3D
     * 
     * @param  D   Datapoints which were fitted
     * @param  fx  fitted plane x values
     * @param  fy  fitted plane y values
     * @param  fy  fitted plane z values
     */
    private void plot3D(double[][] D,double[] fx,double[] fy,double[][] fz)
    {
        Plot3DPanel plot = new Plot3DPanel();
        plot.addScatterPlot("Datapoints", getColumnCopy(D,0), getColumnCopy(D,1), getColumnCopy(D,2));
        plot.addGridPlot("Fitted function", fx, fy, fz);
        JFrame frame = new JFrame("3D Fit Plot");
        frame.setSize(600, 600);
        frame.setContentPane(plot);
        frame.setVisible(true);
    }
}
